import os
import json
import requests
import threading
from time import sleep
import logging


class BaiduMcpClient(object):
    def __init__(self):
        baidu_appkey = os.environ['BAIDUAPPKEY']
        self.mcp_base_url = 'http://appbuilder.baidu.com'
        self.mcp_get_url = f'http://appbuilder.baidu.com/v2/ai_search/mcp/sse?api_key=Bearer+{baidu_appkey}'
        self.mcp_post_url = None
        self.post_id = 0
        self.get_response = None
        self.initialized = False
        self.error = None
        self.tool_call_results = {}
        self.tools = []

        mcp_client_thread = threading.Thread(target=self.mcp_client_loop, args=())
        mcp_client_thread.daemon = True
        mcp_client_thread.start()

        while not self.initialized:
            if self.error is not None:
                logging.critical(self.error)
                raise Exception(self.error)
            sleep(0.1)

    def mcp_client_loop(self):
        while True:  # 一段时间没有访问时对方会挂断连接，因此需要在出现异常后继续循环
            try:
                self.get_response = requests.get(self.mcp_get_url, stream=True)
            except:
                self.error = '连接Baidu MCP SSE Server失败！'
                return

            if self.get_response.status_code != 200:
                self.error = '连接Baidu MCP SSE Server失败！'
                return

            current_event = ''
            try:
                for line in self.get_response.iter_lines():
                    if line:
                        decoded_line = line.decode('utf-8')

                        if decoded_line.startswith(': ping'):
                            continue

                        if decoded_line.startswith('event: endpoint'):
                            current_event = 'endpoint'
                            continue

                        if decoded_line.startswith('event: message'):
                            current_event = 'message'
                            continue

                        if current_event == 'endpoint' and decoded_line.startswith('data: '):
                            endpoint = decoded_line[6:]
                            self.mcp_post_url = self.mcp_base_url + endpoint

                            self.post_id = 1
                            headers = {
                                'Content-Type': 'application/json',
                            }
                            body = {
                                "jsonrpc": "2.0",
                                "id": self.post_id,
                                "method": "initialize",
                                "params": {
                                    "protocolVersion": "2024-11-05",
                                    "capabilities": {
                                        "roots": {
                                            "listChanged": True
                                        },
                                        "sampling": {}
                                    },
                                    "clientInfo": {
                                        "name": "mcp",
                                        "version": "0.1.0"
                                    }
                                }
                            }
                            res = requests.post(self.mcp_post_url, json=body, headers=headers)
                            if res.status_code != 202:
                                self.error = '向Baidu MCP SSE Server发送 initialize 消息失败！'
                                break

                            body = {
                                "jsonrpc": "2.0",
                                "method": "notifications/initialized",
                            }
                            res = requests.post(self.mcp_post_url, json=body, headers=headers)
                            if res.status_code != 202:
                                self.error = '向Baidu MCP SSE Server发送 notifications/initialized 消息失败！'
                                break
                            continue

                        if current_event == 'message' and decoded_line.startswith('data: '):
                            try:
                                message = json.loads(decoded_line[6:])  # 去掉"data: "前缀
                            except:
                                continue

                            if not self.initialized:
                                if message['id'] == 1:
                                    self.initialized = True
                                continue
                            else:
                                if 'result' in message:
                                    if 'content' in message['result']:
                                        if len(message['result']['content']) == 1:
                                            self.tool_call_results[message['id']] = message['result']['content'][0]['text']
                                        else:
                                            self.tool_call_results[message['id']] = message['result']['content']
                                    else:
                                        self.tool_call_results[message['id']] = message['result']
                                else:
                                    self.tool_call_results[message['id']] = message

                                continue
            except Exception as e:
                logging.critical('MCP服务器连接中断，即将重启：' + str(e))

    def call_mcp_tool(self, query):
        self.post_id += 1
        current_post_id = self.post_id

        headers = {
            'Content-Type': 'application/json',
        }
        body = {
            'jsonrpc': '2.0',
            'id': current_post_id,
            'method': 'tools/call',
            'params': {
                'name': 'AIsearch',
                'arguments': {'query': query},
            }
        }

        res = requests.post(self.mcp_post_url, json=body, headers=headers)
        if res.status_code != 202:
            return '调用Baidu MCP服务器工具失败！'

        i = 1
        while True:
            result = self.tool_call_results.get(current_post_id, None)
            if result is not None:
                self.tool_call_results.pop(current_post_id)
                break

            if i > 600:  # 60秒超时
                result = '调用Baidu MCP服务器工具超时！'
                break

            sleep(0.1)
            i += 1

        if not isinstance(result, str):
            result = json.dumps(result, ensure_ascii=False)
        return result

    def close(self):
        self.get_response.close()


if __name__ == "__main__":
    os.environ['BAIDUAPPKEY'] = 'xxxxx'
    client = BaiduMcpClient()

    result = client.call_mcp_tool('人工智能领域最新进展')
    print(result)



